from torch import nn
import torch


class DownSample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DownSample, self).__init__()
        self.conv_relu = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1),
            nn.ReLU()
        )
        self.pool = nn.MaxPool2d(kernel_size=2)

    def forward(self, x, is_pool=True):
        if is_pool:
            x = self.pool(x)
        x = self.conv_relu(x)
        return x


class UpSample(nn.Module):
    def __init__(self, channels):
        super(UpSample, self).__init__()
        self.conv_relu = nn.Sequential(
            nn.Conv2d(in_channels=2 * channels, out_channels=channels, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(in_channels=channels, out_channels=channels, kernel_size=3, padding=1),
            nn.ReLU()
        )
        self.up_conv = nn.Sequential(
            nn.ConvTranspose2d(in_channels=channels, out_channels=channels // 2, kernel_size=3, stride=2,
                               output_padding=1, padding=1),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.conv_relu(x)
        x = self.up_conv(x)
        return x


class UnetModel(nn.Module):
    def __init__(self):
        super(UnetModel, self).__init__()
        self.down_1 = DownSample(in_channels=3, out_channels=64)
        self.down_2 = DownSample(in_channels=64, out_channels=128)
        self.down_3 = DownSample(in_channels=128, out_channels=256)
        self.down_4 = DownSample(in_channels=256, out_channels=512)
        self.down_5 = DownSample(in_channels=512, out_channels=1024)

        self.up = nn.Sequential(
            nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=3, stride=2, output_padding=1,
                               padding=1),
            nn.ReLU()
        )
        self.up_1 = UpSample(channels=512)
        self.up_2 = UpSample(channels=256)
        self.up_3 = UpSample(channels=128)

        self.conv_2 = DownSample(in_channels=128, out_channels=64)
        self.last = nn.Conv2d(in_channels=64, out_channels=2, kernel_size=1)

    def forward(self, x):
        down_1 = self.down_1(x, is_pool=False)
        down_2 = self.down_2(down_1)
        down_3 = self.down_3(down_2)
        down_4 = self.down_4(down_3)
        down_5 = self.down_5(down_4)

        down_5 = self.up(down_5)

        down_5 = torch.cat([down_4, down_5], dim=1)
        down_5 = self.up_1(down_5)

        down_5 = torch.cat([down_3, down_5], dim=1)
        down_5 = self.up_2(down_5)

        down_5 = torch.cat([down_2, down_5], dim=1)
        down_5 = self.up_3(down_5)

        down_5 = torch.cat([down_1, down_5], dim=1)

        down_5 = self.conv_2(down_5, is_pool=False)

        down_5 = self.last(down_5)

        return down_5

